import os
import numpy as np
import torch
import trimesh
from smplx import smplx
from transforms3d.axangles import mat2axangle
from utils.bvh2joint import bvh2joint, default_end_link_trans
from utils.pyt3d_wrapper import Pyt3DWrapper
from utils.visualization import read_data_from_SMPLX, render_SMPLX, save_pcd


def set_theta_from_bvh(smplx_params, id1, id2):
    pass


def check_theta_semantics():
    num_pca_comps = 12
    smplx_model = smplx.create("/share/human_model/models", model_type="smplx", gender="neutral", use_face_contour=False, num_betas=10, num_expression_coeffs=10, ext="npz", use_pca=True, num_pca_comps=num_pca_comps, flat_hand_mean=True)
    smplx_model.to("cuda:0")

    smplx_params = {
        "betas": torch.zeros([1, smplx_model.num_betas], dtype=torch.float32).to("cuda:0"),
        "expression": torch.zeros([1, smplx_model.num_expression_coeffs], dtype=torch.float32).to("cuda:0"),
        "global_orient": torch.zeros([1, 3], dtype=torch.float32).to("cuda:0"),
        "transl": torch.zeros([1, 3], dtype=torch.float32).to("cuda:0"),
        "body_pose": torch.zeros([1, smplx_model.NUM_BODY_JOINTS, 3]).to("cuda:0"),
        "left_hand_pose": torch.zeros([1, num_pca_comps]).to("cuda:0"),
        "right_hand_pose": torch.zeros([1, num_pca_comps]).to("cuda:0"),
    }

    # set something
    data = np.load("/home/liuyun/HHO-dataset/data_processing/exp_data/20230428_data/SMPLX_fitting/1000.npz", allow_pickle=True)["results"].item()
    prep_pose = data["body_pose"]
    print(data["joints"])
    result_model = smplx_model(betas=data["betas"], expression=data["expression"], global_orient=data["global_orient"], transl=data["transl"], body_pose=data["body_pose"], left_hand_pose=data["left_hand_pose"], right_hand_pose=data["right_hand_pose"], return_verts=True)
    result_vertices = result_model.vertices.detach().cpu().numpy()[0]
    result_joints = result_model.joints.detach().cpu().numpy()[0]
    faces = smplx_model.faces_tensor.detach().cpu().numpy()
    print(result_joints)
    # save data
    mesh = trimesh.Trimesh(vertices=result_vertices, faces=faces)
    mesh_txt = trimesh.exchange.obj.export_obj(mesh, include_normals=False, include_color=False, include_texture=False, return_texture=False, write_texture=False, resolver=None, digits=8)
    with open("./ex.obj", "w") as fp:
        fp.write(mesh_txt)
    save_pcd("./ex.ply", result_joints)

    # smplx_params["body_pose"] = prep_pose
    # smplx_params["body_pose"][:, [12,13,15,16,17,18], 0] = 0.
    # smplx_params["body_pose"][:, [0,1,2,3,4,5,8,11,14], 1] = 0.
    # smplx_params["body_pose"][:, [9,10], 2] = 0.

    # set from bvh rotation
    bvh_path = "/home/liuyun/HHO-dataset/data_processing/exp_data/20230323_data/20230323_debug2.bvh"
    end_link_trans = default_end_link_trans()
    _, bvh_rot = bvh2joint(bvh_path, frame_ids=[1029], end_link_trans=end_link_trans, return_local_rot=True)
    T = np.eye(4)
    T[1, 3] = 0.40  # SMPLX TPOSE下pelvis到(0,0,0)的平移
    pelvis2world = bvh_rot[0, 0] @ T
    smplx_params["transl"] = torch.from_numpy(pelvis2world[:3, 3].reshape(1, 3).astype(np.float32)).to("cuda:0")
    d, a = mat2axangle(pelvis2world[:3, :3])
    smplx_params["global_orient"] = torch.from_numpy((d * a).reshape(1, 3).astype(np.float32)).to("cuda:0")
    # set_theta_from_bvh(smplx_params, 0, 1)

    # # render
    # pyt3d_wrapper = Pyt3DWrapper(image_size=(1200, 900), device="cuda:0")
    # render_SMPLX(pyt3d_wrapper, smplx_model, betas=smplx_params["betas"], body_pose=smplx_params["body_pose"], transl=smplx_params["transl"], global_orient=smplx_params["global_orient"], left_hand_pose=smplx_params["left_hand_pose"], right_hand_pose=smplx_params["right_hand_pose"], frame_idx=None)


if __name__ == "__main__":
    check_theta_semantics()
